from src.utils.logger import Logger
from src.utils.config_loader import ConfigLoader
from src.data_processor import DataProcessor
from src.model_manager import ModelManager
from src.evaluator import Evaluator


def main():
    # 初始化日志
    logger = Logger("talent_loss_prediction").get_logger()
    logger.info("====== 模型训练流程启动 ======")

    try:
        # 加载配置
        config = ConfigLoader(r"D:\WorkArea\WorkSpace\Python\talents_loss\config\model_config.yaml")

        # 数据处理
        data_processor = DataProcessor(config, logger)
        X_train, y_train, X_test, y_test = data_processor.load_data()
        preprocessor = data_processor.build_preprocessor(X_train)

        # 任务类型识别
        task_type = data_processor.infer_task_type(y_train)
        logger.info(f"自动识别任务类型: {task_type}")

        # 模型训练
        model_manager = ModelManager(config, logger, task_type)
        top_models = model_manager.train_base_models(X_train, y_train, preprocessor)
        all_models = model_manager.train_stacking_model(X_train, y_train, top_models)

        # 模型评估
        evaluator = Evaluator(config, logger, task_type)
        evaluator.evaluate(X_test, y_test, all_models)

        logger.info("====== 所有流程完成 ======")

    except Exception as e:
        logger.critical(f"程序执行失败: {str(e)}", exc_info=True)
        exit(1)


if __name__ == "__main__":
    main()